feat: RAG follow-ups — kb/query metadata filtering, recursive chunking, Cohere Bedrock embeddings (#153)#171
Conversation
…g, Cohere Bedrock embeddings (#153) Three composable follow-ups on the RAG primitives (#153): - kb/query metadata filtering: an optional `filter` object restricts the search to rows whose metadata contains it via JSONB containment (metadata @> $2::jsonb). The filter binds as a single parameter (no injection surface); a custom metadata_column is supported. Pure prepareQuery core with unit tests (placeholder positioning, custom column, empty-filter omission, non-object/injection errors). - text/chunk recursive mode: `unit: recursive` walks a separator hierarchy (paragraph -> line -> sentence -> word -> character) so chunks break on natural boundaries, then merges pieces up to chunk_size with chunk_overlap characters carried between chunks. Pure, Unicode-aware, unit-tested (sentence boundaries, hard-split fallback, overlap, size bound). rag-ingest now uses it. - Cohere Bedrock embeddings: ai/embed (provider bedrock) adds cohere.embed-english-v3 / cohere.embed-multilingual-v3 alongside Titan. Cohere batches natively (<=96 texts/call) and takes an input_type (search_document default / search_query / classification / clustering), threaded through EmbeddingRequest. Tests cover batching, default/invalid input_type, and order preservation. Docs (RAG guide, connector reference) and examples updated; changeset added.
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
|
Warning Review limit reached
Next review available in: 50 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughThis PR adds three related RAG capabilities to the engine: optional JSONB metadata filtering for ChangesRAG follow-ups
Estimated code review effort: 3 (Moderate) | ~30 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 70af90efaf
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| switch { | ||
| case strings.HasPrefix(req.Model, "amazon.titan-embed"): | ||
| return p.titanEmbeddings(ctx, req) | ||
| case strings.HasPrefix(req.Model, "cohere.embed"): |
There was a problem hiding this comment.
When model is a valid but unsupported Cohere Bedrock ID such as cohere.embed-v4:0, this prefix match routes it into the v3 request path instead of returning the unsupported-model error advertised below. AWS documents Cohere v4 as a separate model with different parameters, including truncate: LEFT|RIGHT|NONE, so the hard-coded v3 body with truncate: "END" can fail at runtime or return embeddings with dimensions this connector does not handle; restrict this branch to the two v3 IDs the connector actually supports.
Useful? React with 👍 / 👎.
The broad HasPrefix(model, "cohere.embed") match routed any Cohere ID — including newer models like cohere.embed-v4:0, which use a different request/response shape (truncate LEFT|RIGHT|NONE, different output) — into the v3 code path with a hard-coded v3 body. Match only the two v3 families (cohere.embed-english-v3 / -multilingual-v3) so unsupported IDs fall through to the unsupported-model error. Adds a regression test for v4.
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
packages/engine/internal/connector/kb.go (1)
299-327: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMisleading
nextParamsemantics inkbFilterClause.The doc comment says
nextParamis "the 1-based index of the next free placeholder ($1 is always the query vector)", but the caller passeslen(args)(which is1for the vector-only case) and the function itself adds+1to compute the actual placeholder ($2). SonextParamis really "the count of already-bound params," not "the next free index" as documented. The result is correct, but the naming/comment could confuse future maintainers into passing an off-by-one value.✏️ Suggested clarification
-// An absent or empty filter yields no clause. nextParam is the 1-based index of -// the next free placeholder ($1 is always the query vector). -func kbFilterClause(params map[string]any, nextParam int) (kbFilter, error) { +// An absent or empty filter yields no clause. boundParamCount is the number of +// SQL parameters already bound (currently always 1, for the query vector); +// the filter placeholder is boundParamCount+1. +func kbFilterClause(params map[string]any, boundParamCount int) (kbFilter, error) {And update the usage of
nextParam+1toboundParamCount+1accordingly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/engine/internal/connector/kb.go` around lines 299 - 327, The `kbFilterClause` comment and parameter name are misleading because the caller passes the count of already-bound args, not the “next free placeholder” index; clarify this by renaming `nextParam` to `boundParamCount` (or equivalent) in `kbFilterClause` and updating the placeholder expression to use that meaning consistently. Also update the doc comment to say it’s the number of already-bound parameters, since `$1` is reserved for the query vector and the filter should bind at `boundParamCount + 1`. Keep the behavior unchanged, but make the naming and comment match the actual usage in the `kbFilterClause` helper.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/engine/internal/connector/chunk.go`:
- Around line 121-151: The mergeSplits chunking contract is inconsistent with
the current behavior in mergeSplits, since flushing then retaining overlap can
produce chunks larger than size; update the implementation to either enforce the
max before appending or adjust the documented/tests contract to reflect the real
size + overlap bound. Check the mergeSplits function and any related chunk-size
docstrings or expectations in connector/chunk.go so the code and comments agree.
In `@packages/site/src/content/docs/workflow-reference/connectors.md`:
- Around line 407-412: The Markdownlint MD038 failure comes from an inline code
span in the chunking strategy description using leading or trailing whitespace.
Inspect the affected markdown in the connectors documentation around the `unit`
field and the `recursive` chunking description, then trim any spaces inside
backticks so every inline code span is flush with its content.
---
Nitpick comments:
In `@packages/engine/internal/connector/kb.go`:
- Around line 299-327: The `kbFilterClause` comment and parameter name are
misleading because the caller passes the count of already-bound args, not the
“next free placeholder” index; clarify this by renaming `nextParam` to
`boundParamCount` (or equivalent) in `kbFilterClause` and updating the
placeholder expression to use that meaning consistently. Also update the doc
comment to say it’s the number of already-bound parameters, since `$1` is
reserved for the query vector and the filter should bind at `boundParamCount +
1`. Keep the behavior unchanged, but make the naming and comment match the
actual usage in the `kbFilterClause` helper.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 958abe18-dff7-4e9c-a606-960a45c65d59
📒 Files selected for processing (13)
.changeset/rag-followups.mdpackages/engine/internal/connector/chunk.gopackages/engine/internal/connector/chunk_test.gopackages/engine/internal/connector/embed.gopackages/engine/internal/connector/kb.gopackages/engine/internal/connector/kb_test.gopackages/engine/internal/connector/provider.gopackages/engine/internal/connector/provider_bedrock.gopackages/engine/internal/connector/provider_bedrock_embed_test.gopackages/site/src/content/docs/rag-guide.mdpackages/site/src/content/docs/workflow-reference/connectors.mdpackages/site/src/content/examples/rag-ask.yamlpackages/site/src/content/examples/rag-ingest.yaml
…Clause naming - chunk.go: mergeSplits/chunkText docs now state a recursive chunk targets `size` and may reach `size + overlap` (carried overlap tail), matching the size-bound test and the connector reference — not "at most size". - connectors.md: quote the recursive separators (`". "`, `" "`, etc.) so no inline code span has a leading/trailing space (markdownlint MD038). - kb.go: rename kbFilterClause's `nextParam` to `boundParamCount` and clarify the doc — it's the count of already-bound params ($1 = vector), so the filter binds at boundParamCount+1. Behavior unchanged.
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
2 participants
Three composable follow-ups on the RAG primitives from #153.
Changes
kb/querymetadata filtering — an optionalfilterobject scopes the vector search to rows whose metadata contains it, via the JSONB containment operator (metadata @> $2::jsonb). The whole filter binds as a single parameter (no injection surface); a custommetadata_columnis supported, and an empty object is ignored. PureprepareQuerycore.recursivechunking —text/chunkgainsunit: recursive, which walks a separator hierarchy (paragraph → line → sentence → word → character) so chunks break on natural boundaries instead of mid-sentence, then merges pieces up tochunk_sizewithchunk_overlapcharacters carried between chunks. Unicode-aware.rag-ingestnow uses it.ai/embed(providerbedrock) addscohere.embed-english-v3andcohere.embed-multilingual-v3alongside Amazon Titan. Cohere batches natively (up to 96 texts per call) and takes aninput_type(search_documentdefault /search_query/classification/clustering), threaded throughEmbeddingRequest. Cohere's Bedrock response carries no token counts, so those embeddings report zero usage.Docs (RAG guide, connector reference), both RAG examples, and a changeset are updated.
Testing
prepareQueryfilter (placeholder positioning, custom column, empty-filter omission, non-object/injection rejection); recursive chunking (sentence boundaries, hard-split fallback, overlap, size bound); Cohere embeddings (batching, default/invalid input_type, order preservation).go build ./cmd/mantle,go vet ./..., gofmt (changed files) clean.mantle validatepasses onrag-ingest.yamlandrag-ask.yaml.Related Issues
Follow-ups on #153. Token-accurate (tokenizer-based) chunking remains the one open item there.
🤖 Generated with Claude Code
Generated by Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Documentation